import cv2
import numpy as np
import ncnn
import time

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -88.72, 88.72)))

def preprocess(img, target_size=320):
    h, w = img.shape[:2]
    scale = min(target_size / w, target_size / h)
    new_w, new_h = int(w * scale), int(h * scale)
    resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    top = (target_size - new_h) // 2
    bottom = target_size - new_h - top
    left = (target_size - new_w) // 2
    right = target_size - new_w - left
    padded = cv2.copyMakeBorder(resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
    blob = padded.astype(np.float32) / 255.0
    blob = np.transpose(blob, (2, 0, 1))  # HWC to CHW
    return blob, scale, (left, top)

def postprocess(preds, scale, pad, conf_thres=0.25, iou_thres=0.45, target_size=320):
    num_classes = 3  # 实际类别数为 3
    if preds.shape[1] != 4 + num_classes:
        raise ValueError(f"Expected {4 + num_classes} columns, got {preds.shape[1]}")

    bboxes = preds[:, :4]      # [cx, cy, w, h] in PIXELS (0~320 range)
    cls_scores = preds[:, 4:]  # raw logits for 3 classes

    # Apply sigmoid to class scores only
    cls_scores = sigmoid(cls_scores)
    max_scores = np.max(cls_scores, axis=1)
    class_ids = np.argmax(cls_scores, axis=1)

    # Filter by confidence
    keep = max_scores >= conf_thres
    if not np.any(keep):
        return [], [], [], []

    bboxes = bboxes[keep]
    scores = max_scores[keep]
    classes = class_ids[keep]

    # Convert [cx, cy, w, h] (pixels in 320x320) → [x1, y1, x2, y2]
    cx = bboxes[:, 0]
    cy = bboxes[:, 1]
    w = bboxes[:, 2]
    h = bboxes[:, 3]

    x1 = cx - w / 2
    y1 = cy - h / 2
    x2 = cx + w / 2
    y2 = cy + h / 2

    # Remove padding and scale back to original image
    x1 = (x1 - pad[0]) / scale
    y1 = (y1 - pad[1]) / scale
    x2 = (x2 - pad[0]) / scale
    y2 = (y2 - pad[1]) / scale

    boxes = np.stack([x1, y1, x2, y2], axis=1).astype(np.int32)

    # NMS
    indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), conf_thres, iou_thres)
    if len(indices) == 0:
        return [], [], [], []
    if isinstance(indices[0], list):  # OpenCV 4.7+
        indices = [i[0] for i in indices]

    return boxes[indices], scores[indices], classes[indices], cls_scores[keep][indices]

def draw_boxes(img, boxes, scores, classes):
    class_names = ["Class0", "Class1", "Class2"]  # 可替换为实际类别名
    for (x1, y1, x2, y2), score, cls in zip(boxes, scores, classes):
        color = [(0, 255, 0), (255, 0, 0), (0, 0, 255)][cls % 3]
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
        label = f"{class_names[cls]}: {score:.2f}"
        cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
    return img

# -------------------------------
# Main
# -------------------------------
if __name__ == "__main__":
    param_path = r"D:\NCNN\ncnn-20230816-windows-vs2017\x64\bin\202510161157\best31n.param"
    bin_path = r"D:\NCNN\ncnn-20230816-windows-vs2017\x64\bin\202510161157\best31n.bin"
    img_path = r"C:\TestPic\test-onnx\detect_20250927_103320_1758940400649.jpg"
    output_path = f"C:\\TestPic\\test-onnx\\predict\\ncnn_result_{int(time.time())}.jpg"

    img = cv2.imread(img_path)
    if img is None:
        raise FileNotFoundError(f"Image not found: {img_path}")

    blob, scale, pad = preprocess(img, target_size=320)

    net = ncnn.Net()
    # net.opt.use_vulkan_compute = True  # 可选：启用 GPU
    # net.opt.use_int8_inference = True  # 若已量化 INT8
    net.load_param(param_path)
    net.load_model(bin_path)

    with net.create_extractor() as ex:
        ex.input("images", ncnn.Mat(blob))
        ret, out = ex.extract("output0")  # 输出层名来自 .param 文件

    if ret != 0:
        raise RuntimeError("NCNN inference failed")

    # 🔧 关键修正：直接 reshape 为 (2100, 7)
    raw = np.array(out)
    if raw.size != 2100 * 7:
        raise ValueError(f"Unexpected number of elements: {raw.size}, expected {2100 * 7}")
    pred = raw.reshape((2100, 7))  # ✅ 直接指定逻辑 shape

    print(f"out.w={out.w}, out.h={out.h}, out.c={out.c}")
    print(f"Raw output shape: {pred.shape}")  # 应为 (2100, 7)
    print(f"Sample prediction: {pred[0]}")    # 值应在 0~320 范围（像素坐标）

    boxes, scores, classes, _ = postprocess(
        pred, scale, pad, conf_thres=0.5, iou_thres=0.45, target_size=320
    )

    print(f"Detected {len(boxes)} boxes")
    result_img = draw_boxes(img.copy(), boxes, scores, classes)
    cv2.imwrite(output_path, result_img)
    cv2.imshow("NCNN YOLOv8 Prediction", result_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()